# Multi-Species Support for FLEXINVERT.jl
# Extending the framework to handle multiple trace gases simultaneously

module MultiSpecies

using ..CoreTypes: Domain, State, Observations
using ..Covariance: CovarianceMatrix, build_prior_covariance
using ..ForwardModel: forward_model, calculate_model_concentrations
using ..Settings: Config, PriorFluxes
using LinearAlgebra
using Statistics
using Printf

export MultiSpeciesConfig, MultiSpeciesState, MultiSpeciesObservations, MultiSpeciesFluxes,
       run_multi_species_inversion, SpeciesDefinition, create_correlation_matrix

"""
    SpeciesDefinition

Definition of a trace gas species with its properties.

# Fields
- `name::Symbol`: Species name (e.g., :CO2, :CH4, :N2O)
- `molar_mass::Float64`: Molar mass in g/mol
- `background_lifetime::Float64`: Background atmospheric lifetime (years)
- `measurement_precision::Float64`: Typical measurement precision (ppb)
- `units::String`: Concentration units (ppm, ppb, ppt)
- `conversion_factor::Float64`: Conversion factor to standard units
- `color::String`: Plot color for visualization
- `priority::Int`: Processing priority (lower = higher priority)
"""
Base.@kwdef struct SpeciesDefinition
    name::Symbol
    molar_mass::Float64
    background_lifetime::Float64
    measurement_precision::Float64
    units::String
    conversion_factor::Float64 = 1.0
    color::String = "blue"
    priority::Int = 1
end

# Predefined species definitions
const SPECIES_DEFINITIONS = Dict{Symbol, SpeciesDefinition}(
    :CO2 => SpeciesDefinition(
        name = :CO2,
        molar_mass = 44.01,
        background_lifetime = 100.0,  # Very long-lived
        measurement_precision = 0.1,   # 0.1 ppm
        units = "ppm",
        conversion_factor = 1.0,
        color = "red",
        priority = 1
    ),
    :CH4 => SpeciesDefinition(
        name = :CH4,
        molar_mass = 16.04,
        background_lifetime = 12.0,
        measurement_precision = 2.0,   # 2 ppb
        units = "ppb",
        conversion_factor = 1000.0,    # ppb to ppm
        color = "green",
        priority = 2
    ),
    :N2O => SpeciesDefinition(
        name = :N2O,
        molar_mass = 44.01,
        background_lifetime = 114.0,
        measurement_precision = 0.3,   # 0.3 ppb
        units = "ppb",
        conversion_factor = 1000.0,    # ppb to ppm
        color = "orange",
        priority = 3
    ),
    :CO => SpeciesDefinition(
        name = :CO,
        molar_mass = 28.01,
        background_lifetime = 2.0,
        measurement_precision = 10.0,  # 10 ppb
        units = "ppb",
        conversion_factor = 1000.0,    # ppb to ppm
        color = "purple",
        priority = 4
    ),
    :H2 => SpeciesDefinition(
        name = :H2,
        molar_mass = 2.016,
        background_lifetime = 2.0,
        measurement_precision = 10.0,  # 10 ppb
        units = "ppb",
        conversion_factor = 1000.0,    # ppb to ppm
        color = "brown",
        priority = 5
    )
)

"""
    MultiSpeciesConfig

Configuration for multi-species inversion.

# Fields
- `species::Vector{Symbol}`: List of species to invert
- `species_configs::Dict{Symbol, Config}`: Individual species configurations
- `correlation_matrix::Matrix{Float64}`: Cross-species correlation matrix
- `joint_inversion::Bool`: Whether to perform joint inversion or separate inversions
- `correlation_length::Float64`: Spatial correlation length for cross-species correlations
- `temporal_correlation::Float64`: Temporal correlation for cross-species
- `regularization::Float64`: Regularization strength for joint inversion
- `use_species_priors::Bool`: Whether to use species-specific priors
"""
Base.@kwdef mutable struct MultiSpeciesConfig
    # Species selection
    species::Vector{Symbol} = [:CO2, :CH4]
    species_configs::Dict{Symbol, Config} = Dict()

    # Cross-species correlations
    correlation_matrix::Matrix{Float64} = Matrix{Float64}(I, 2, 2)
    correlation_length::Float64 = 500.0  # km
    temporal_correlation::Float64 = 7.0   # days

    # Inversion settings
    joint_inversion::Bool = true
    regularization::Float64 = 1e-6
    use_species_priors::Bool = true

    # Performance settings
    parallel_species::Bool = true
    memory_efficient::Bool = false

    # Output settings
    save_cross_covariances::Bool = true
    compute_correlations::Bool = true
end

"""
    MultiSpeciesObservations

Container for multi-species observation data.

# Fields
- `observations::Dict{Symbol, Observations}`: Individual species observations
- `shared_stations::Vector{String}`: Stations with measurements for multiple species
- `time_alignment::Dict{Symbol, Float64}`: Time alignment offsets for each species
- `measurement_correlations::Dict{Tuple{Symbol, Symbol}, Float64}`: Measurement error correlations
- `total_observations::Int`: Total number of observations across all species
"""
mutable struct MultiSpeciesObservations
    observations::Dict{Symbol, Observations}
    shared_stations::Vector{String}
    time_alignment::Dict{Symbol, Float64}
    measurement_correlations::Dict{Tuple{Symbol, Symbol}, Float64}
    total_observations::Int

    function MultiSpeciesObservations()
        return new(
            Dict{Symbol, Observations}(),
            String[],
            Dict{Symbol, Float64}(),
            Dict{Tuple{Symbol, Symbol}, Float64}(),
            0
        )
    end
end

"""
    MultiSpeciesFluxes

Container for multi-species prior fluxes.

# Fields
- `fluxes::Dict{Symbol, PriorFluxes}`: Individual species fluxes
- `cross_fluxes::Dict{Tuple{Symbol, Symbol}, Array{Float64, 3}}`: Cross-species flux interactions
- `emission_factors::Dict{Tuple{Symbol, Symbol}, Float64}`: Emission conversion factors
- `lifetime_corrections::Dict{Symbol, Float64}`: Lifetime-based corrections
"""
struct MultiSpeciesFluxes
    fluxes::Dict{Symbol, PriorFluxes}
    cross_fluxes::Dict{Tuple{Symbol, Symbol}, Array{Float64, 3}}
    emission_factors::Dict{Tuple{Symbol, Symbol}, Float64}
    lifetime_corrections::Dict{Symbol, Float64}

    function MultiSpeciesFluxes()
        return new(
            Dict{Symbol, PriorFluxes}(),
            Dict{Tuple{Symbol, Symbol}, Array{Float64, 3}}(),
            Dict{Tuple{Symbol, Symbol}, Float64}(),
            Dict{Symbol, Float64}()
        )
    end
end

"""
    MultiSpeciesState

State representation for multi-species inversion.

# Fields
- `species_states::Dict{Symbol, State}`: Individual species states
- `joint_state::Union{State, Nothing}`: Joint state (if joint inversion)
- `cross_covariances::Dict{Tuple{Symbol, Symbol}, Matrix{Float64}}`: Cross-species covariances
- `species_correlations::Matrix{Float64}`: Current species correlation estimates
- `joint_cost::Float64`: Joint cost function value
"""
mutable struct MultiSpeciesState
    species_states::Dict{Symbol, State}
    joint_state::Union{State, Nothing}
    cross_covariances::Dict{Tuple{Symbol, Symbol}, Matrix{Float64}}
    species_correlations::Matrix{Float64}
    joint_cost::Float64

    function MultiSpeciesState(species::Vector{Symbol})
        return new(
            Dict{Symbol, State}(),
            nothing,
            Dict{Tuple{Symbol, Symbol}, Matrix{Float64}}(),
            Matrix{Float64}(I, length(species), length(species)),
            0.0
        )
    end
end

"""
    create_correlation_matrix(species::Vector{Symbol}, base_correlations::Dict{Tuple{Symbol, Symbol}, Float64})

Create cross-species correlation matrix.
"""
function create_correlation_matrix(species::Vector{Symbol}, base_correlations::Dict{Tuple{Symbol, Symbol}, Float64} = Dict{Tuple{Symbol, Symbol}, Float64}())
    n_species = length(species)
    corr_matrix = Matrix{Float64}(I, n_species, n_species)

    # Fill in known correlations
    for (species_pair, correlation) in base_correlations
        sp1, sp2 = species_pair
        if sp1 in species && sp2 in species
            idx1 = findfirst(==(sp1), species)
            idx2 = findfirst(==(sp2), species)
            if idx1 !== nothing && idx2 !== nothing
                corr_matrix[idx1, idx2] = correlation
                corr_matrix[idx2, idx1] = correlation
            end
        end
    end

    # Add reasonable default correlations for related species
    # CO2-CH4: negative correlation (fossil fuel sources)
    if :CO2 in species && :CH4 in species
        idx_co2 = findfirst(==(:CO2), species)
        idx_ch4 = findfirst(==(:CH4), species)
        if idx_co2 !== nothing && idx_ch4 !== nothing && corr_matrix[idx_co2, idx_ch4] == 0.0
            corr_matrix[idx_co2, idx_ch4] = -0.3
            corr_matrix[idx_ch4, idx_co2] = -0.3
        end
    end

    # CO-CO2: positive correlation (combustion sources)
    if :CO in species && :CO2 in species
        idx_co = findfirst(==(:CO), species)
        idx_co2 = findfirst(==(:CO2), species)
        if idx_co !== nothing && idx_co2 !== nothing && corr_matrix[idx_co, idx_co2] == 0.0
            corr_matrix[idx_co, idx_co2] = 0.5
            corr_matrix[idx_co2, idx_co] = 0.5
        end
    end

    return corr_matrix
end

"""
    build_multi_species_covariance(config::MultiSpeciesConfig, domain::Domain, fluxes::MultiSpeciesFluxes)

Build joint covariance matrix for multi-species inversion.
"""
function build_multi_species_covariance(config::MultiSpeciesConfig, domain::Domain, fluxes::MultiSpeciesFluxes)
    species = config.species
    n_species = length(species)

    @info "Building multi-species covariance matrix for $(n_species) species"

    # Build individual species covariances
    individual_covs = Dict{Symbol, CovarianceMatrix}()
    total_dims = 0

    for sp in species
        if haskey(fluxes.fluxes, sp) && haskey(config.species_configs, sp)
            sp_cov = build_prior_covariance(config.species_configs[sp], domain, fluxes.fluxes[sp])
            individual_covs[sp] = sp_cov
            total_dims += sp_cov.n_modes
        end
    end

    if total_dims == 0
        error("No valid species covariances found")
    end

    # Build joint covariance matrix
    joint_cov = zeros(Float64, total_dims, total_dims)
    species_indices = Dict{Symbol, UnitRange{Int}}()
    current_idx = 1

    # Fill diagonal blocks with individual covariances
    for sp in species
        if haskey(individual_covs, sp)
            sp_cov = individual_covs[sp]
            sp_dim = sp_cov.n_modes
            sp_range = current_idx:(current_idx + sp_dim - 1)

            joint_cov[sp_range, sp_range] = sp_cov.full
            species_indices[sp] = sp_range
            current_idx += sp_dim
        end
    end

    # Fill off-diagonal blocks with cross-species correlations
    if config.joint_inversion
        for (i, sp1) in enumerate(species)
            for (j, sp2) in enumerate(species)
                if i < j && haskey(species_indices, sp1) && haskey(species_indices, sp2)
                    idx1 = species_indices[sp1]
                    idx2 = species_indices[sp2]
                    correlation = config.correlation_matrix[i, j]

                    # Simple correlation model: scale individual covariances
                    cov1_sqrt = sqrt.(abs.(diag(joint_cov[idx1, idx1])))
                    cov2_sqrt = sqrt.(abs.(diag(joint_cov[idx2, idx2])))

                    cross_cov = correlation * (cov1_sqrt * cov2_sqrt')
                    joint_cov[idx1, idx2] = cross_cov
                    joint_cov[idx2, idx1] = cross_cov'
                end
            end
        end
    end

    # Apply regularization
    if config.regularization > 0
        joint_cov += I * config.regularization
    end

    @info "Built joint covariance matrix: $(size(joint_cov))"

    vals, vecs = eigen(Symmetric(joint_cov))
    perm = sortperm(vals, rev=true)
    eigenvalues = vals[perm]
    eigenvectors = vecs[:, perm]
    sqrt_eigs = sqrt.(eigenvalues)
    inv_sqrt_eigs = 1.0 ./ sqrt_eigs
    inv_eigs = 1.0 ./ eigenvalues

    identity1 = Matrix{Float64}(I, 1, 1)
    ones_vec = ones(Float64, total_dims)
    cov = CovarianceMatrix(
        joint_cov,
        eigenvalues,
        eigenvectors,
        identity1,
        identity1,
        identity1,
        identity1,
        length(eigenvalues),
        sum(eigenvalues),
        1.0,
        1, 1, total_dims, total_dims,
        ones_vec,
        ones_vec,
        false,
        ones_vec,
        ones_vec,
        false,
        ones_vec,
        ones_vec,
        false,
        ones_vec,
        ones_vec,
        false,
        nothing,
        1,
        1,
        nothing,
        nothing,
        nothing,
        nothing,
        nothing,
        nothing,
        0.0,
        true
    )

    return cov, individual_covs, species_indices
end

"""
    run_multi_species_inversion(observations::MultiSpeciesObservations, domain::Domain,
                                fluxes::MultiSpeciesFluxes, config::MultiSpeciesConfig)

Run multi-species inversion.
"""
function run_multi_species_inversion(observations::MultiSpeciesObservations, domain::Domain,
                                    fluxes::MultiSpeciesFluxes, config::MultiSpeciesConfig)

    @info "Starting multi-species inversion for: $(config.species)"

    if config.joint_inversion
        return run_joint_inversion(observations, domain, fluxes, config)
    else
        return run_separate_inversions(observations, domain, fluxes, config)
    end
end

"""
    run_joint_inversion(observations, domain, fluxes, config)

Run joint multi-species inversion.
"""
function run_joint_inversion(observations::MultiSpeciesObservations, domain::Domain,
                            fluxes::MultiSpeciesFluxes, config::MultiSpeciesConfig)

    species = config.species
    n_species = length(species)

    @info "Running joint inversion for $n_species species"

    # Build joint covariance matrix
    joint_cov, individual_covs, species_indices = build_multi_species_covariance(config, domain, fluxes)

    # Create joint observation vector
    joint_observations = concatenate_observations(observations, species)

    # Build joint observation operator
    joint_H = build_joint_observation_operator(observations, domain, fluxes, species_indices)

    # Initialize joint state
    joint_state = initialize_joint_state(species, individual_covs, species_indices)

    # Solve joint inversion
    result = solve_joint_inversion(joint_state, joint_observations, domain, joint_cov, joint_H, config)

    # Split results back to individual species
    species_results = split_joint_results(result, species, species_indices)

    @info "Joint inversion completed successfully"

    return species_results, result
end

"""
    run_separate_inversions(observations, domain, fluxes, config)

Run separate inversions for each species.
"""
function run_separate_inversions(observations::MultiSpeciesObservations, domain::Domain,
                               fluxes::MultiSpeciesFluxes, config::MultiSpeciesConfig)

    species = config.species
    species_results = Dict{Symbol, State}()

    @info "Running separate inversions for $(length(species)) species"

    for sp in species
        @info "Processing species: $sp"

        if haskey(observations.observations, sp) && haskey(fluxes.fluxes, sp)
            sp_obs = observations.observations[sp]
            sp_fluxes = fluxes.fluxes[sp]
            sp_config = get(config.species_configs, sp, Config())

            # Run single-species inversion
            sp_result = run_single_species_inversion(sp_obs, domain, sp_fluxes, sp_config)
            species_results[sp] = sp_result

            @info "Completed inversion for $sp"
        else
            @warn "Missing observations or fluxes for species $sp"
        end
    end

    @info "Separate inversions completed"

    return species_results, nothing
end

"""
    concatenate_observations(observations::MultiSpeciesObservations, species::Vector{Symbol})

Concatenate observations from multiple species into joint observation vector.
"""
function concatenate_observations(observations::MultiSpeciesObservations, species::Vector{Symbol})
    total_observations = 0
    observation_data = []

    # Count total observations
    for sp in species
        if haskey(observations.observations, sp)
            sp_obs = observations.observations[sp]
            total_observations += length(sp_obs.concentrations)
        end
    end

    @info "Concatenating $total_observations total observations"

    # Build joint observation vector
    joint_concentrations = zeros(total_observations)
    joint_errors = zeros(total_observations)
    joint_stations = String[]
    joint_datetimes = DateTime[]

    current_idx = 1
    for sp in species
        if haskey(observations.observations, sp)
            sp_obs = observations.observations[sp]
            n_sp_obs = length(sp_obs.concentrations)

            # Convert to standard units (ppm)
            sp_def = SPECIES_DEFINITIONS[sp]
            conversion = sp_def.conversion_factor

            joint_concentrations[current_idx:current_idx+n_sp_obs-1] = sp_obs.concentrations / conversion
            joint_errors[current_idx:current_idx+n_sp_obs-1] = sp_obs.measurement_errors / conversion
            append!(joint_stations, sp_obs.stations)
            append!(joint_datetimes, sp_obs.datetimes)

            current_idx += n_sp_obs
        end
    end

    return (
        concentrations = joint_concentrations,
        errors = joint_errors,
        stations = joint_stations,
        datetimes = joint_datetimes,
        n_observations = total_observations
    )
end

"""
    build_joint_observation_operator(observations, domain, fluxes, species_indices)

Build joint observation operator for multi-species inversion.
"""
function build_joint_observation_operator(observations::MultiSpeciesObservations, domain::Domain,
                                        fluxes::MultiSpeciesFluxes, species_indices::Dict{Symbol, UnitRange{Int}})

    species = collect(keys(species_indices))
    total_dims = sum(length(species_indices[sp]) for sp in species)

    # Count total observations
    total_obs = 0
    for sp in species
        if haskey(observations.observations, sp)
            total_obs += length(observations.observations[sp].concentrations)
        end
    end

    @info "Building joint observation operator: $(total_obs) × $(total_dims)"

    # Initialize joint H matrix
    joint_H = zeros(Float64, total_obs, total_dims)

    # Fill in blocks for each species
    obs_row_start = 1
    for sp in species
        if haskey(observations.observations, sp) && haskey(species_indices, sp)
            sp_obs = observations.observations[sp]
            sp_def = SPECIES_DEFINITIONS[sp]
            n_sp_obs = length(sp_obs.concentrations)

            # Convert units
            conversion = sp_def.conversion_factor

            # Extract relevant columns
            sp_cols = species_indices[sp]
            sp_obs_row_end = obs_row_start + n_sp_obs - 1

            # Fill block (simplified - would use actual footprint data)
            joint_H[obs_row_start:sp_obs_row_end, sp_cols] = build_species_observation_operator(
                sp_obs, domain, fluxes.fluxes[sp], conversion
            )

            obs_row_start = sp_obs_row_end + 1
        end
    end

    return joint_H
end

"""
    build_species_observation_operator(obs, domain, fluxes, conversion)

Build observation operator for a single species.
"""
function build_species_observation_operator(obs::Observations, domain::Domain, fluxes::PriorFluxes, conversion::Float64)
    n_obs = length(obs.concentrations)
    n_spatial = domain.nx * domain.ny

    # Simplified implementation - would use actual footprints
    H = zeros(Float64, n_obs, n_spatial)

    # Fill with transport model (placeholder)
    for i in 1:n_obs
        # Each observation responds to nearby grid cells
        H[i, :] = 0.1 * rand(n_spatial) / conversion
    end

    return H
end

"""
    initialize_joint_state(species, individual_covs, species_indices)

Initialize joint state vector.
"""
function initialize_joint_state(species::Vector{Symbol}, individual_covs::Dict{Symbol, CovarianceMatrix},
                             species_indices::Dict{Symbol, UnitRange{Int}})

    total_dims = sum(length(species_indices[sp]) for sp in species if haskey(individual_covs, sp))
    joint_control = zeros(total_dims)
    joint_physical = zeros(total_dims)

    # Initialize each species block
    for sp in species
        if haskey(individual_covs, sp) && haskey(species_indices, sp)
            sp_idx = species_indices[sp]
            n_sp = length(sp_idx)

            # Initialize with prior mean (zero)
            joint_control[sp_idx] = zeros(n_sp)
            joint_physical[sp_idx] = zeros(n_sp)
        end
    end

    return (
        control = joint_control,
        physical = joint_physical,
        total_dims = total_dims
    )
end

"""
    solve_joint_inversion(state, observations, domain, covariance, H, config)

Solve joint inversion problem.
"""
function solve_joint_inversion(state, observations, domain, covariance, H, config)
    n_obs, n_control = size(H)

    @info "Solving joint inversion: $n_obs observations, $n_control control variables"

    # Build observation error covariance
    R = diagm(observations.errors.^2)

    # Solve using normal equations
    A = H' * inv(R) * H + inv(covariance.full)
    b = H' * inv(R) * observations.concentrations

    # Add regularization
    if config.regularization > 0
        A += I * config.regularization
    end

    # Solve
    try
        solution = cholesky(A) \ b
        @info "Joint inversion solved using Cholesky decomposition"
    catch e
        @warn "Cholesky failed, using SVD fallback: $e"
        U, S, V = svd(A)
        solution = V * diag(1 ./ S) * U' * b
    end

    # Compute cost components
    model_concentrations = H * solution
    residuals = observations.concentrations - model_concentrations

    J_obs = 0.5 * residuals' * inv(R) * residuals
    J_prior = 0.5 * solution' * inv(covariance.full) * solution
    J_total = J_obs + J_prior

    @info "Joint inversion completed:"
    @info "  J_obs = $J_obs"
    @info "  J_prior = $J_prior"
    @info "  J_total = $J_total"

    return (
        solution = solution,
        model_concentrations = model_concentrations,
        residuals = residuals,
        J_obs = J_obs,
        J_prior = J_prior,
        J_total = J_total,
        convergence = true
    )
end

"""
    split_joint_results(result, species, species_indices)

Split joint inversion results back to individual species.
"""
function split_joint_results(result, species::Vector{Symbol}, species_indices::Dict{Symbol, UnitRange{Int}})
    species_results = Dict{Symbol, State}()

    for sp in species
        if haskey(species_indices, sp)
            sp_idx = species_indices[sp]
            sp_solution = result.solution[sp_idx]

            # Create species state
            sp_state = State()
            sp_state.control = sp_solution
            sp_state.physical = sp_solution  # Simplified transformation
            sp_state.J_obs = result.J_obs / length(species)  # Distribute cost
            sp_state.J_prior = result.J_prior / length(species)
            sp_state.J_total = result.J_total / length(species)

            species_results[sp] = sp_state
        end
    end

    return species_results
end

"""
    run_single_species_inversion(obs, domain, fluxes, config)

Run inversion for a single species (used in separate inversion mode).
"""
function run_single_species_inversion(obs::Observations, domain::Domain, fluxes::PriorFluxes, config::Config)
    # This would call the existing single-species inversion code
    # For now, return a simple result
    state = State()
    state.control = zeros(100)  # Placeholder
    state.physical = zeros(100)
    state.J_obs = 10.0
    state.J_prior = 5.0
    state.J_total = 15.0

    return state
end

"""
    compute_multi_species_diagnostics(results::Dict{Symbol, State}, config::MultiSpeciesConfig)

Compute diagnostics for multi-species inversion results.
"""
function compute_multi_species_diagnostics(results::Dict{Symbol, State}, config::MultiSpeciesConfig)
    diagnostics = Dict()

    species = config.species

    # Individual species diagnostics
    for sp in species
        if haskey(results, sp)
            sp_result = results[sp]
            diagnostics[sp] = Dict(
                :J_obs => sp_result.J_obs,
                :J_prior => sp_result.J_prior,
                :J_total => sp_result.J_total,
                :n_control => length(sp_result.control),
                :solution_norm => norm(sp_result.control)
            )
        end
    end

    # Cross-species correlations
    if length(species) > 1
        correlation_matrix = zeros(length(species), length(species))
        for (i, sp1) in enumerate(species)
            for (j, sp2) in enumerate(species)
                if i <= j && haskey(results, sp1) && haskey(results, sp2)
                    sol1 = results[sp1].control
                    sol2 = results[sp2].control

                    if i == j
                        correlation_matrix[i, j] = 1.0
                    else
                        # Compute correlation between species solutions
                        min_len = min(length(sol1), length(sol2))
                        if min_len > 1
                            corr = cor(sol1[1:min_len], sol2[1:min_len])
                            correlation_matrix[i, j] = corr
                            correlation_matrix[j, i] = corr
                        end
                    end
                end
            end
        end
        diagnostics[:cross_correlations] = correlation_matrix
    end

    return diagnostics
end

end # module
